package com.cg.ai.rag;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.io.File;
import java.util.List;

/**
 * RAG (检索增强生成) 配置
 * 加载知识库文档并提供内容检索功能
 */
@Configuration
@Slf4j
public class RagConfig {

    @Autowired
    private EmbeddingModel qwenEmbeddingModel;

    /**
     * 创建内存向量存储
     */
    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }

    /**
     * 创建内容检索器
     * @return 配置好的ContentRetriever实例
     */
    @Bean
    public ContentRetriever contentRetriever() {
        try {
            EmbeddingStore<TextSegment> embeddingStore = embeddingStore();
            
            // 从classpath加载知识库文档目录
            String docsPath = null;
            File docsDir = null;
            
            try {
                var resource = getClass().getClassLoader().getResource("docs");
                if (resource != null) {
                    docsDir = new File(resource.toURI());
                    docsPath = docsDir.getAbsolutePath();
                    log.info("从classpath加载知识库文档目录: {}", docsPath);
                } else {
                    log.info("classpath中未找到docs目录，RAG功能将使用空的向量存储");
                }
            } catch (Exception e) {
                log.warn("从classpath加载知识库文档目录失败: {}", e.getMessage());
            }
            
            if (docsDir != null && docsDir.exists() && docsDir.isDirectory()) {
                // 1. 加载文档
                List<Document> documents = FileSystemDocumentLoader.loadDocuments(docsPath);
                log.info("加载了 {} 个知识库文档", documents.size());
                
                if (!documents.isEmpty()) {
                    // 2. 文档切割：将每个文档按每段进行分割，最大 1000 字符，每次重叠最多 200 个字符
                    DocumentByParagraphSplitter paragraphSplitter = new DocumentByParagraphSplitter(1000, 200);
                    
                    // 3. 自定义文档加载器
                    EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                            .documentSplitter(paragraphSplitter)
                            // 为了提高搜索质量，为每个 TextSegment 添加文档名称
                            .textSegmentTransformer(textSegment -> TextSegment.from(
                                    textSegment.metadata().getString("file_name") + "\n" + textSegment.text(),
                                    textSegment.metadata()
                            ))
                            // 使用指定的向量模型
                            .embeddingModel(qwenEmbeddingModel)
                            .embeddingStore(embeddingStore)
                            .build();
                    
                    // 加载文档到向量存储
                    ingestor.ingest(documents);
                    log.info("知识库文档已成功加载到向量存储");
                }
            } else {
                log.info("知识库文档目录未找到或为空，RAG功能将使用空的向量存储");
            }
            
            // 4. 创建内容查询器
            return EmbeddingStoreContentRetriever.builder()
                    .embeddingStore(embeddingStore)
                    .embeddingModel(qwenEmbeddingModel)
                    .maxResults(5) // 最多 5 个检索结果
                    .minScore(0.75) // 过滤掉分数小于 0.75 的结果
                    .build();
                    
        } catch (Exception e) {
            log.error("RAG配置初始化失败", e);
            // 返回一个空的检索器，避免启动失败
            return EmbeddingStoreContentRetriever.builder()
                    .embeddingStore(new InMemoryEmbeddingStore<>())
                    .embeddingModel(qwenEmbeddingModel)
                    .maxResults(1)
                    .minScore(0.0)
                    .build();
        }
    }
}

